import re
from typing import Dict, Iterator, Optional
from collections import defaultdict

from chernc.treediff.gumtree.tree.default_tree import DefaultTree
from chernc.treediff.gumtree.tree.tree import TreeLike


class TreeContext:
    """
    The tree context class contains an AST together with its context (key - value metadata).
    """

    def __init__(self):
        self.metadata: Dict[str, object] = {}
        self.serializers = self.MetadataSerializers()
        self.root: Optional[TreeLike] = None

    def __str__(self):
        return str(TreeIoUtils.toText(self))

    def set_root(self, root: TreeLike):
        """
        Set the AST of the TreeContext.
        """
        self.root = root

    def get_root(self) -> Optional[TreeLike]:
        """
        Return the AST of the TreeContext.
        """
        return self.root

    def create_tree(self, type_: str, label: str) -> TreeLike:
        """
        Utility method to create a default tree with a given type and label.
        """
        return DefaultTree(type_, label)

    def create_tree_without_label(self, type_: str) -> TreeLike:
        """
        Utility method to create a default tree with a given type.
        """
        return DefaultTree(type_)

    def create_fake_tree(self, *trees: TreeLike) -> TreeLike:
        """
        Utility method to create a fake tree with the given children.
        """
        return FakeTree(trees)

    def get_metadata(self, key: str) -> Optional[object]:
        """
        Get the AST metadata with the given key.
        There is no way to know if the metadata is really null or does not exist.
        """
        return self.metadata.get(key)

    def set_metadata(self, key: str, value: object) -> Optional[object]:
        """
        Store an AST metadata with the given key and value.
        """
        return self.metadata.setdefault(key, value)

    def get_metadata_iterator(self) -> Iterator:
        """
        Get an iterator on the AST metadata.
        """
        return iter(self.metadata.items())

    def get_serializers(self) -> 'MetadataSerializers':
        """
        Get the metadata serializers for this tree context.
        """
        return self.serializers

    def export(self, serializers: 'MetadataSerializers') -> 'TreeContext':
        """
        Export metadata serializers.
        """
        self.serializers.add_all(serializers)
        return self

    def export_single(self, key: str, serializer: 'MetadataSerializer') -> 'TreeContext':
        """
        Export a single metadata serializer.
        """
        self.serializers.add(key, serializer)
        return self

    def export_multiple(self, *names: str) -> 'TreeContext':
        """
        Export multiple metadata serializers with default string conversion.
        """
        for name in names:
            self.serializers.add(name, lambda x: str(x))
        return self

    class Marshallers:
        """
        Handles adding and removing serializers with valid keys.
        """
        valid_id = re.compile(r"[a-zA-Z0-9_]*")

        def __init__(self):
            self.serializers: Dict[str, object] = {}

        def add_all(self, other: 'Marshallers') -> None:
            """
            Add all serializers from another Marshallers instance.
            """
            self.add_all_from_dict(other.serializers)

        def add_all_from_dict(self, serializers: Dict[str, object]) -> None:
            """
            Add all serializers from a dictionary.
            """
            for key, serializer in serializers.items():
                self.add(key, serializer)

        def add(self, name: str, serializer: object) -> None:
            """
            Add a single serializer.
            """
            if not self.valid_id.match(name):
                raise RuntimeError("Invalid key for serialization")
            self.serializers[name] = serializer

        def remove(self, key: str) -> None:
            """
            Remove a serializer by key.
            """
            if key in self.serializers:
                del self.serializers[key]

        def exports(self) -> set:
            """
            Get the set of exported serializer names.
            """
            return set(self.serializers.keys())

    class MetadataSerializers(Marshallers):
        """
        Handles the serialization of metadata for the TreeContext.
        """

        def serialize(self, formatter: 'TreeFormatter', key: str, value: object) -> None:
            """
            Serialize metadata to a formatter.
            """
            serializer = self.serializers.get(key)
            if serializer:
                formatter.serialize_attribute(key, serializer.to_string(value))

    class MetadataUnserializers(Marshallers):
        """
        Handles the deserialization of metadata for the TreeContext.
        """

        def load(self, tree: TreeLike, key: str, value: str) -> None:
            """
            Load and unserialize metadata into a tree.
            """
            unserializer = self.serializers.get(key)
            if unserializer:
                if key == "pos":
                    tree.set_pos(int(value))
                elif key == "length":
                    tree.set_length(int(value))
                else:
                    tree.set_metadata(key, unserializer.from_string(value))
